package org.raymond.iworks.gateway.filter.support;

import cn.hutool.core.codec.Base64;
import cn.hutool.crypto.symmetric.AES;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.raymond.iworks.apis.common.ErrorCode;
import org.raymond.iworks.apis.common.Result;
import org.raymond.iworks.apis.exception.BaseException;
import org.raymond.iworks.commons.utils.ErrorUtil;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.factory.rewrite.RewriteFunction;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

@Slf4j
public class EncryptFunction implements RewriteFunction<String, String> {
    private final boolean encrypt;
    private final ObjectMapper objectMapper;

    public EncryptFunction(boolean encrypt, ObjectMapper objectMapper){
        this.encrypt = encrypt;
        this.objectMapper = objectMapper;
    }

    @Override
    public Publisher<String> apply(ServerWebExchange exchange, String responseString) {
        log.info("responseString:{}", responseString);
        String responseBody = "";
        GatewayContext gatewayContext = exchange.getAttribute(GatewayContext.CACHE_GATEWAY_CONTEXT);
        if (gatewayContext == null) {
            throw new BaseException(ErrorCode.ERROR, "Gateway context should not be null");
        }
        AES aes = gatewayContext.getAes();
        try {
            Result result = objectMapper.readValue(responseString, Result.class);
            String resultStr = result.getData()==null?"":objectMapper.writeValueAsString(result.getData());
            String encBase64 = "";
            if(encrypt && aes!=null){
                encBase64 = aes.encryptBase64(resultStr);
            }else{
                encBase64 = Base64.encode(resultStr);
            }
            log.info("str:{}, b64:{}", resultStr, encBase64);
            result.setData(encBase64);
            responseBody = objectMapper.writeValueAsString(result);
        } catch (Exception e) {
            log.error("Write response body error: {}", ErrorUtil.getErrorLog(e));
        }
        return Mono.just(responseBody);
    }
}
